package msu.ml.data;

import msu.ml.core.*;
import msu.ml.util.*;
import weka.core.*;
import java.io.*;
import java.util.*;

/**
 * The CrossValidationDataProvider divides the data or a subset
 * of the data in a directory into a number of folds for use in
 * a cross validation experiment. Using cross validation, if there
 * are n folds, then n-1 folds are used for training and the final 
 * fold is used for validation. This cycle repeats n times, thus all
 * files are classified exactly once.
 *
 * @author Reginald M Mead
 * @version 1.0
 */
public class DbcvDataProvider implements IDataProvider
{
   private String [][] folds;
   private int filesPerFold;
   private int numFolds;
   private ArrayList<IDataPreProcessor> filters;
   private ArrayList<String> volumeNames;
   private DatabaseModel db;
   private int currentFold;

   /**
    * Creates a new CrossValidationDataProvider
    *
    */
   public DbcvDataProvider(String url, String user, String pass)
   {
      this.db = new DatabaseModel(url, user, pass);
      this.filesPerFold = 0;
      this.filters = new ArrayList<IDataPreProcessor>();
      this.volumeNames = new ArrayList<String>();
   }

   public void addFilter(IDataPreProcessor filter)
   {
      this.filters.add(filter);
   }

   public void clearFilters()
   {
      this.filters.clear();
   }

   public void reset()
   {
      this.currentFold = 0;
   }

   public int count()
   {
      if(!this.isInitialized)
         initialize();

      return this.numFolds * this.filesPerFold;
   }

   public void addVolume(String volumeName)
   {
      this.volumeNames.add(volumeName);
   }

   /**
    * Provides the next group of target data instances.
    * Unlike the training data which can be aggregated into
    * one monolythic instances object, the target data
    * is provides as an array of instances, each representing
    * a single sweep. In this way the data can be classified
    * one sweep at a time.
    *
    * @return next group of target data instances
    */
   public NxInstances [] getNextTargetData()
   {
      if(!this.isInitialized)
         initialize();

      NxInstances [] result = new NxInstances[this.filesPerFold];
      for(int i = 0; i < this.filesPerFold; i++)
      {
         result[i] = new NxInstances(this.db.getInstances(
            folds[this.currentFold][i]));
         for(int k = 0; k < filters.size(); k++)
            result[i] = filters.get(k).process(result[i]);
      }

      this.currentFold++;

      return result;
   }

   /**
    * Provides the next set of training instances
    *
    * @return next training data
    */
   public NxInstances [] getNextIndividualTrainingData()
   {
      if(!this.isInitialized)
         initialize();

      NxInstances [] result = new NxInstances[this.filesPerFold * (this.numFolds - 1)];
      for(int i = 0, k = 0; i < this.numFolds; i++)
      {
         if(i != this.currentFold)
         {
            for(int j = 0; j < folds[i].length; j++, k++)
            {
               result[k] = new NxInstances(this.db.getInstances(folds[i][j]));
               for(int f = 0; f < filters.size(); f++)
                  result[k] = filters.get(f).process(result[k]);
            }
         }
      }

      return result;
   }

   /**
    * Provides the next set of training instances in one
    * large Instances object.
    *
    * @return next training data
    */
   public NxInstances getNextTrainingData()
   {
      if(!this.isInitialized)
         initialize();

      NxInstances result = null;
      for(int i = 0; i < this.numFolds; i++)
      {
         if(i != this.currentFold)
         {
            for(int j = 0; j < folds[i].length; j++)
            {
               if(result == null)
                  result = new NxInstances(this.db.getInstances(folds[i][j]));
               else
                  result.concatenate(this.db.getInstances(folds[i][j]));
            }
         }
      }

      for(int i = 0; i < filters.size(); i++)
         result = filters.get(i).process(result);

      return result;
   }

   /**
    * Get the names of the data sources used in the
    * next set of training data.
    *
    * @return the labels of the next training data
    */
   public String [] getNextTrainingLabels()
   {
      if(!this.isInitialized)
         initialize();

      String [] result = new String[(this.numFolds - 1) * this.filesPerFold];
      for(int i = 0, k = 0; i < this.numFolds; i++)
         if(i != this.currentFold)
            for(int j = 0; j < folds[i].length; j++)
               result[k++] = folds[i][j];

      return result;
   }

   /**
    * return the number of folds to divide the data into.
    *
    * @return number of folds
    */
   public int getNumFolds()
   {
      return numFolds;
   }

   /**
    * Return whether there is more data 
    * @return true if more data exists
    */
   public boolean hasMoreData()
   {
      if(!this.isInitialized)
         initialize();

      return this.currentFold < this.numFolds;
   }

   /**
    * Print a listing of the files in a given fold
    *
    * @param index index of fold to print
    */
   public void printFold(int index)
   {
      System.out.println();
      System.out.println("Fold #" + index);
      for(int i = 0; i < folds[index].length; i++)
      {
         System.out.println(folds[index][i]);
      }
   }

   public String [] getFileNames()
   {
      String [] result = new String[this.numFolds * this.filesPerFold];
      for(int i = 0, idx = 0; i < this.numFolds; i++)
         for(int j = 0; j < this.filesPerFold; j++)
            result[idx++] = folds[i][j];
      return result;
   }

   /**
    * set the number of folds to use
    *
    * @param newVal number of folds to use
    */
   public void setNumFolds(int newVal)
   {
      numFolds = newVal;
   }

   protected boolean isInitialized = false;
   public void initialize()
   {
      List<String> volumes = null;

      if(this.volumeNames.size() > 0)
         volumes = new ArrayList<String>(this.volumeNames);
      else
         volumes = db.getVolumeNames();

      this.filesPerFold = volumes.size() / this.numFolds;

      HashMap<String, Integer> classes = db.getClassesForVolumes(volumes);

      ArrayList<String> class0 = new ArrayList<String>();
      ArrayList<String> class1 = new ArrayList<String>();
      for(int i = 0; i < volumes.size(); i++)
      {
         if(classes.get(volumes.get(i)).intValue() == 0)
            class0.add(volumes.get(i));
         else if(classes.get(volumes.get(i)).intValue() == 1)
            class1.add(volumes.get(i));
      }

      while(class0.size() > (this.numFolds * this.filesPerFold / 2))
         class0.remove((int)(Math.random() * (double)class0.size()));
      while(class1.size() > (this.numFolds * this.filesPerFold / 2))
         class1.remove((int)(Math.random() * (double)class1.size()));
      while(class0.size() < (this.numFolds * this.filesPerFold / 2))
         class0.add(class0.get((int)(Math.random() * (double)class0.size())));
      while(class1.size() < (this.numFolds * this.filesPerFold / 2))
         class1.add(class1.get((int)(Math.random() * (double)class1.size())));

      System.out.println("Non-Biological Sweeps: ");
      for(int i = 0; i < class0.size(); i++)
         System.out.println(class0.get(i));
      System.out.println("Biological Sweeps: ");
      for(int i = 0; i < class1.size(); i++)
         System.out.println(class1.get(i));

      volumes.clear();
      volumes.addAll(class0);
      volumes.addAll(class1);

      this.filesPerFold = volumes.size() / this.numFolds;
      this.folds = new String[this.numFolds][this.filesPerFold];

      //Construct Folds
      for(int i = 0; i < this.numFolds; i++)
      {
         for(int j = 0; j < folds[i].length; j++)
         {
            int index = (int)(Math.random() * (double)volumes.size());
            folds[i][j] = volumes.remove(index);
         }
      }

      this.currentFold = 0;
      this.isInitialized = true;
   }
}
